//
//  MFAPMNetMonitorURLSession.m
//  MFAPMManager
//
//  Created by zhang liyu on 2019/6/27.
//

#import "MFAPMNetMonitorURLSession.h"

#import "AFHTTPSessionManager+MFAPMNetMonitor.h"

typedef void (^AFURLSessionTaskProgressBlock)(NSProgress *);

@interface MFAPMNetMonitorURLSessionTaskInfo : NSObject

@property (atomic, readonly, copy) NSArray *modes;
@property (atomic, readonly, strong) NSThread *thread;
@property (atomic, readonly, strong) NSURLSessionDataTask *task;
@property (atomic, readonly, strong) id<NSURLSessionDataDelegate>  delegate;
@property (nonatomic, copy) AFURLSessionTaskProgressBlock uploadProgressBlock;
@property (nonatomic, copy) AFURLSessionTaskProgressBlock downloadProgressBlock;

- (instancetype)initWithTask:(NSURLSessionDataTask *)task delegate:(id< NSURLSessionDataDelegate>)delegate modes:(NSArray *)modes;

- (void)performBlock:(dispatch_block_t)block;

- (void)invalidate;

- (void)cleanUpProgressForTask:(NSURLSessionTask *)task;

@end

@interface MFAPMNetMonitorURLSessionTaskInfo ()

@property (atomic, readwrite, strong) id<NSURLSessionDataDelegate>  delegate;
@property (atomic, readwrite, strong) NSThread *thread;

@property (nonatomic, strong) NSProgress *uploadProgress;
@property (nonatomic, strong) NSProgress *downloadProgress;

//kvo监听的函数字符串
@property (nonatomic, copy) NSString *countOfReceive;
@property (nonatomic, copy) NSString *countOfExpectReceive;
@property (nonatomic, copy) NSString *countOfSend;
@property (nonatomic, copy) NSString *countOfExpectSend;

@end

@implementation MFAPMNetMonitorURLSessionTaskInfo

#pragma mark - Life Circle

- (instancetype)initWithTask:(NSURLSessionDataTask *)task delegate:(id<NSURLSessionDataDelegate>)delegate modes:(NSArray *)modes{
    self = [super init];
    if (self != nil) {
        self->_task = task;
        self->_delegate = delegate;
        self->_thread = [NSThread currentThread];
        self->_modes = [modes copy];
        
        //初始化上传下载Progress
        self->_uploadProgress = [[NSProgress alloc] initWithParent:nil userInfo:nil];
        self->_uploadProgress.totalUnitCount = NSURLSessionTransferSizeUnknown;
        
        self->_downloadProgress = [[NSProgress alloc] initWithParent:nil userInfo:nil];
        self->_downloadProgress.totalUnitCount = NSURLSessionTransferSizeUnknown;
        
        //为Task添加进度监听
        [self setupProgressForTask:self->_task];
    }
    return self;
}

#pragma mark - NSProgress Tracking

- (void)setupProgressForTask:(NSURLSessionTask *)task {
    __weak __typeof__(task) weakTask = task;
    
    self.uploadProgress.totalUnitCount = task.countOfBytesExpectedToSend;
    self.downloadProgress.totalUnitCount = task.countOfBytesExpectedToReceive;
    [self.uploadProgress setCancellable:YES];
    [self.uploadProgress setCancellationHandler:^{
        __typeof__(weakTask) strongTask = weakTask;
        [strongTask cancel];
    }];
    [self.uploadProgress setPausable:YES];
    [self.uploadProgress setPausingHandler:^{
        __typeof__(weakTask) strongTask = weakTask;
        [strongTask suspend];
    }];
    if ([self.uploadProgress respondsToSelector:@selector(setResumingHandler:)]) {
        [self.uploadProgress setResumingHandler:^{
            __typeof__(weakTask) strongTask = weakTask;
            [strongTask resume];
        }];
    }
    
    [self.downloadProgress setCancellable:YES];
    [self.downloadProgress setCancellationHandler:^{
        __typeof__(weakTask) strongTask = weakTask;
        [strongTask cancel];
    }];
    [self.downloadProgress setPausable:YES];
    [self.downloadProgress setPausingHandler:^{
        __typeof__(weakTask) strongTask = weakTask;
        [strongTask suspend];
    }];
    
    if ([self.downloadProgress respondsToSelector:@selector(setResumingHandler:)]) {
        [self.downloadProgress setResumingHandler:^{
            __typeof__(weakTask) strongTask = weakTask;
            [strongTask resume];
        }];
    }
    
    [task addObserver:self
           forKeyPath:NSStringFromSelector(@selector(countOfBytesReceived))
              options:NSKeyValueObservingOptionNew
              context:NULL];
    [task addObserver:self
           forKeyPath:NSStringFromSelector(@selector(countOfBytesExpectedToReceive))
              options:NSKeyValueObservingOptionNew
              context:NULL];
    
    [task addObserver:self
           forKeyPath:NSStringFromSelector(@selector(countOfBytesSent))
              options:NSKeyValueObservingOptionNew
              context:NULL];
    [task addObserver:self
           forKeyPath:NSStringFromSelector(@selector(countOfBytesExpectedToSend))
              options:NSKeyValueObservingOptionNew
              context:NULL];
    
    [self.downloadProgress addObserver:self
                            forKeyPath:NSStringFromSelector(@selector(fractionCompleted))
                               options:NSKeyValueObservingOptionNew
                               context:NULL];
    [self.uploadProgress addObserver:self
                          forKeyPath:NSStringFromSelector(@selector(fractionCompleted))
                             options:NSKeyValueObservingOptionNew
                             context:NULL];
}

- (void)cleanUpProgressForTask:(NSURLSessionTask *)task {
    [task removeObserver:self forKeyPath:NSStringFromSelector(@selector(countOfBytesReceived))];
    [task removeObserver:self forKeyPath:NSStringFromSelector(@selector(countOfBytesExpectedToReceive))];
    [task removeObserver:self forKeyPath:NSStringFromSelector(@selector(countOfBytesSent))];
    [task removeObserver:self forKeyPath:NSStringFromSelector(@selector(countOfBytesExpectedToSend))];
    [self.downloadProgress removeObserver:self forKeyPath:NSStringFromSelector(@selector(fractionCompleted))];
    [self.uploadProgress removeObserver:self forKeyPath:NSStringFromSelector(@selector(fractionCompleted))];
}

- (void)observeValueForKeyPath:(NSString *)keyPath ofObject:(id)object change:(NSDictionary<NSString *,id> *)change context:(void *)context {
    if ([object isKindOfClass:[NSURLSessionTask class]] || [object isKindOfClass:[NSURLSessionDownloadTask class]]) {
        if ([keyPath isEqualToString:self.countOfReceive]) {
            self.downloadProgress.completedUnitCount = [change[NSKeyValueChangeNewKey] longLongValue];
        } else if ([keyPath isEqualToString:self.countOfExpectReceive]) {
            self.downloadProgress.totalUnitCount = [change[NSKeyValueChangeNewKey] longLongValue];
        } else if ([keyPath isEqualToString:self.countOfSend]) {
            self.uploadProgress.completedUnitCount = [change[NSKeyValueChangeNewKey] longLongValue];
        } else if ([keyPath isEqualToString:self.countOfExpectSend]) {
            self.uploadProgress.totalUnitCount = [change[NSKeyValueChangeNewKey] longLongValue];
        }
    }
    else if ([object isEqual:self.downloadProgress]) {
        if (self->_downloadProgressBlock) {
            __weak __typeof__(self) weakSelf = self;
            [self performBlock:^{
                weakSelf.downloadProgressBlock(object);
            }];
        }
    }
    else if ([object isEqual:self.uploadProgress]) {
        if (self->_uploadProgressBlock) {
            __weak __typeof__(self) weakSelf = self;
            [self performBlock:^{
                weakSelf.uploadProgressBlock(object);
            }];
        }
    }
}

#pragma mark - Block Perform Actions

- (void)performBlock:(dispatch_block_t)block{
    [self performSelector:@selector(performBlockOnClientThread:) onThread:self.thread withObject:[block copy] waitUntilDone:NO modes:self.modes];
}

- (void)performBlockOnClientThread:(dispatch_block_t)block{
    block();
}

- (void)invalidate{
    self.delegate = nil;
    self.thread = nil;
}

#pragma mark - Setter & Getter

- (NSString *)countOfReceive{
    if(!_countOfReceive){
        _countOfReceive = NSStringFromSelector(@selector(countOfBytesReceived));
    }
    return _countOfReceive;
}

- (NSString *)countOfExpectReceive{
    if(!_countOfExpectReceive){
        _countOfExpectReceive = NSStringFromSelector(@selector(countOfBytesExpectedToReceive));
    }
    return _countOfExpectReceive;
}

- (NSString *)countOfSend{
    if(!_countOfSend){
        _countOfSend = NSStringFromSelector(@selector(countOfBytesSent));
    }
    return _countOfSend;
}

- (NSString *)countOfExpectSend{
    if(!_countOfExpectSend){
        _countOfExpectSend = NSStringFromSelector(@selector(countOfBytesExpectedToSend));
    }
    return _countOfExpectSend;
}

@end


@interface MFAPMNetMonitorURLSession () <NSURLSessionDataDelegate>

@property (atomic, readonly, strong) NSMutableDictionary *taskInfoByTaskID;
@property (atomic, readonly, strong) NSOperationQueue *sessionDelegateQueue;

@end

@implementation MFAPMNetMonitorURLSession

#pragma mark - Install Actions

- (instancetype)init{
    return [self initWithConfiguration:nil];
}

- (instancetype)initWithConfiguration:(nullable NSURLSessionConfiguration *)configuration{
    self = [super init];
    if (self != nil) {
        if (configuration == nil) {
            configuration = [NSURLSessionConfiguration defaultSessionConfiguration];
        }
        self->_configuration = [configuration copy];
        
        self->_taskInfoByTaskID = [[NSMutableDictionary alloc] init];
        
        self->_sessionDelegateQueue = [[NSOperationQueue alloc] init];
        [self->_sessionDelegateQueue setMaxConcurrentOperationCount:1];
        [self->_sessionDelegateQueue setName:@"cn.MissFresh.NetMonitor.Session.Queue"];
        
        self->_session = [NSURLSession sessionWithConfiguration:self->_configuration delegate:self delegateQueue:self->_sessionDelegateQueue];
        self->_session.sessionDescription = @"cn.MissFresh.NetMonitor.Session.Queue";
        
        //交换AFHTTPSessionManager的方法，用于获取progress回调
        [AFHTTPSessionManager mfapm_swizzleRequestProgressActions];
    }
    return self;
}

#pragma mark - Task Actions

- (NSURLSessionDataTask *)dataTaskWithRequest:(NSURLRequest *)request delegate:(id<NSURLSessionDataDelegate>)delegate modes:(NSArray *)modes{
    
    if ([modes count] == 0) {
        modes = @[NSDefaultRunLoopMode];
    }
    
    NSString *urlStr = request.URL.absoluteString;
    
    NSURLSessionDataTask *task = [self.session dataTaskWithRequest:request];
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [[MFAPMNetMonitorURLSessionTaskInfo alloc] initWithTask:task delegate:delegate modes:modes];
    taskInfo.uploadProgressBlock = [[AFHTTPSessionManager mfapm_uploadProgressBKForKey:urlStr] copy];
    taskInfo.downloadProgressBlock = [[AFHTTPSessionManager mfapm_downloadProgressBKForKey:urlStr] copy];

    @synchronized (self) {
        self.taskInfoByTaskID[@(task.taskIdentifier)] = taskInfo;
    }
    return task;
}

- (MFAPMNetMonitorURLSessionTaskInfo *)taskInfoForTask:(NSURLSessionTask *)task{
    
    MFAPMNetMonitorURLSessionTaskInfo *result;
    @synchronized (self) {
        result = self.taskInfoByTaskID[@(task.taskIdentifier)];
    }
    return result;
}

#pragma mark - Session Data Delegate Actions

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task willPerformHTTPRedirection:(NSHTTPURLResponse *)response newRequest:(NSURLRequest *)newRequest completionHandler:(void (^)(NSURLRequest *))completionHandler{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:task];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:task:willPerformHTTPRedirection:newRequest:completionHandler:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session task:task willPerformHTTPRedirection:response newRequest:newRequest completionHandler:completionHandler];
        }];
    } else {
        completionHandler(newRequest);
    }
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didReceiveChallenge:(NSURLAuthenticationChallenge *)challenge completionHandler:(void (^)(NSURLSessionAuthChallengeDisposition disposition, NSURLCredential *credential))completionHandler{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:task];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:task:didReceiveChallenge:completionHandler:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session task:task didReceiveChallenge:challenge completionHandler:completionHandler];
        }];
    } else {
        completionHandler(NSURLSessionAuthChallengePerformDefaultHandling, nil);
    }
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task needNewBodyStream:(void (^)(NSInputStream *bodyStream))completionHandler{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:task];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:task:needNewBodyStream:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session task:task needNewBodyStream:completionHandler];
        }];
    } else {
        completionHandler(nil);
    }
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didSendBodyData:(int64_t)bytesSent totalBytesSent:(int64_t)totalBytesSent totalBytesExpectedToSend:(int64_t)totalBytesExpectedToSend{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:task];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:task:didSendBodyData:totalBytesSent:totalBytesExpectedToSend:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session task:task didSendBodyData:bytesSent totalBytesSent:totalBytesSent totalBytesExpectedToSend:totalBytesExpectedToSend];
        }];
    }
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didCompleteWithError:(NSError *)error{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:task];
    // This is our last delegate callback so we remove our task info record.
    
    @synchronized (self) {
        [taskInfo cleanUpProgressForTask:taskInfo.task];
        [self.taskInfoByTaskID removeObjectForKey:@(taskInfo.task.taskIdentifier)];
    }
    
    // Call the delegate if required.  In that case we invalidate the task info on the client thread
    // after calling the delegate, otherwise the client thread side of the -performBlock: code can
    // find itself with an invalidated task info.
    
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:task:didCompleteWithError:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session task:task didCompleteWithError:error];
            [taskInfo invalidate];
        }];
    } else {
        [taskInfo invalidate];
    }
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveResponse:(NSURLResponse *)response completionHandler:(void (^)(NSURLSessionResponseDisposition disposition))completionHandler{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:dataTask];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:dataTask:didReceiveResponse:completionHandler:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session dataTask:dataTask didReceiveResponse:response completionHandler:completionHandler];
        }];
    } else {
        completionHandler(NSURLSessionResponseAllow);
    }
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didBecomeDownloadTask:(NSURLSessionDownloadTask *)downloadTask{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:dataTask];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:dataTask:didBecomeDownloadTask:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session dataTask:dataTask didBecomeDownloadTask:downloadTask];
        }];
    }
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveData:(NSData *)data{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:dataTask];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:dataTask:didReceiveData:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session dataTask:dataTask didReceiveData:data];
        }];
    }
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask willCacheResponse:(NSCachedURLResponse *)proposedResponse completionHandler:(void (^)(NSCachedURLResponse *cachedResponse))completionHandler{
    
    MFAPMNetMonitorURLSessionTaskInfo *taskInfo = [self taskInfoForTask:dataTask];
    if ([taskInfo.delegate respondsToSelector:@selector(URLSession:dataTask:willCacheResponse:completionHandler:)]) {
        [taskInfo performBlock:^{
            [taskInfo.delegate URLSession:session dataTask:dataTask willCacheResponse:proposedResponse completionHandler:completionHandler];
        }];
    } else {
        completionHandler(proposedResponse);
    }
}

@end
